import sys
import cv2
import numpy as np
from PyQt5 import QtCore
from PyQt5.QtWidgets import QApplication, QMainWindow, QPushButton, QVBoxLayout, QHBoxLayout, QWidget, QLabel, QTextEdit, QMessageBox
from PyQt5.QtCore import QThread, pyqtSignal, Qt
from PyQt5.QtGui import QImage, QPixmap, QFont
from paddleocr import PaddleOCR
from PIL import Image, ImageDraw, ImageFont

class VideoThread(QThread):
    change_pixmap_signal = pyqtSignal(QImage)
    update_text_signal = pyqtSignal(str)

    def __init__(self, camera_index=0):
        super().__init__()
        self.camera_index = camera_index
        self.run_flag = False
        self.ocr = PaddleOCR(use_angle_cls=True, lang="ch")
        self.cap = cv2.VideoCapture(camera_index)

    def run(self):
        self.run_flag = True
        while self.run_flag:
            ret, frame = self.cap.read()
            if ret:
                result, frame = self.ocr_process(frame)
                if result:
                    self.update_text_signal.emit(result)
                self.change_pixmap_signal.emit(QImage(frame, frame.shape[1], frame.shape[0], QImage.Format_RGB888).rgbSwapped())
        self.cap.release()

    def stop(self):
        self.run_flag = False
        self.wait()

    def ocr_process(self, frame):
        result = self.ocr.ocr(frame, cls=True)
        if not result[0]:
            return "", frame
        else:
            pil_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
            draw = ImageDraw.Draw(pil_img)
            font = ImageFont.truetype('simsun.ttc', 12)
            final_text = ""
            for box, text in zip([x[0] for x in result[0]], [x[1][0] for x in result[0]]):
                x = [point[0] for point in box]
                y = [point[1] for point in box]
                x0, y0 = min(x), min(y)
                x1, y1 = max(x), max(y)
                draw.rectangle(((x0, y0), (x1, y1)), outline=(0, 255, 0), width=2)
                draw.text((x0, y0 - 20), text, font=font, fill=(0, 255, 0))
                final_text += text + "\n"
            return final_text, cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("OCR For Video")
        self.setFixedSize(960, 540)
        self.setWindowFlags(Qt.Window)
        self.camera_index = 0
        self.initUI()
        self.thread = None

    def initUI(self):
        self.central_widget = QWidget(self)
        self.setCentralWidget(self.central_widget)
        self.main_layout = QVBoxLayout(self.central_widget)

        self.banner = QWidget(self)
        # self.banner.setStyleSheet("border: 1px solid black; background-color: #ccc;")
        self.banner_layout = QHBoxLayout(self.banner)
        self.banner_layout.setContentsMargins(10, 10, 10, 10)  # 设置边距
        self.banner_layout.setSpacing(10)  # 设置间距

        self.logo_label = QLabel(self)
        self.logo_label.setPixmap(QPixmap("./logo.png").scaled(80, 80))  # 设置本地图片并缩放
        self.logo_label.setFixedSize(80, 80)  # 设置图片大小
        self.banner_layout.addWidget(self.logo_label)

        self.company_label = QLabel("牛叉科技", self)
        # self.company_label.setStyleSheet("font: bold 1px;")
        self.company_label.setStyleSheet("font: bold 18px; border: 1px solid black; background-color: #ccc;")
        self.banner_layout.addWidget(self.company_label)

        self.team_label = QLabel("码海逸帆客", self)
        # self.team_label.setStyleSheet("font: 18px;")
        self.team_label.setStyleSheet("font: 18px; border: 1px solid black; background-color: #ccc;")
        self.banner_layout.addWidget(self.team_label)

        self.main_layout.addWidget(self.banner)

        self.splitter = QHBoxLayout()
        self.main_layout.addLayout(self.splitter)

        self.video_frame = QLabel(self)
        self.video_frame.setStyleSheet("border: 1px solid black; background-color: #ccc;")
        self.video_frame.setFixedSize(600, 400)

        self.video_frame.setGeometry(QtCore.QRect(0, 0, 600, 400))
        self.video_frame.setMinimumSize(QtCore.QSize(600, 400))
        self.video_frame.setMaximumSize(QtCore.QSize(600, 400))


        self.splitter.addWidget(self.video_frame)

        self.controls_area = QWidget(self)
        self.splitter.addWidget(self.controls_area)
        controls_layout = QVBoxLayout(self.controls_area)

        button_layout = QHBoxLayout()
        controls_layout.addLayout(button_layout)

        self.start_btn = QPushButton("启动", self)
        self.start_btn.setStyleSheet("background-color: lightgreen;")
        self.start_btn.setFixedSize(80, 40)
        button_layout.addWidget(self.start_btn)

        self.stop_btn = QPushButton("停止", self)
        self.stop_btn.setStyleSheet("background-color: lightblue;")
        self.stop_btn.setFixedSize(80, 40)
        button_layout.addWidget(self.stop_btn)

        self.save_btn = QPushButton("保存", self)
        self.save_btn.setStyleSheet("background-color: lightcoral;")
        self.save_btn.setFixedSize(80, 40)
        button_layout.addWidget(self.save_btn)

        self.start_btn.clicked.connect(self.start_video)
        self.stop_btn.clicked.connect(self.stop_video)
        self.save_btn.clicked.connect(self.save_text)

        self.text_area = QTextEdit(self)
        self.text_area.setFont(QFont("SimSun", 16))
        self.text_area.setStyleSheet("border: 1px solid black; background-color: #999;")
        self.text_area.setReadOnly(True)
        controls_layout.addWidget(self.text_area)

        self.thread = VideoThread(self.camera_index)
        self.thread.change_pixmap_signal.connect(self.update_image)
        self.thread.update_text_signal.connect(self.update_text)

    def start_video(self):
        if self.thread is None or not self.thread.isRunning():
            self.thread = VideoThread(self.camera_index)
            self.thread.change_pixmap_signal.connect(self.update_image)
            self.thread.update_text_signal.connect(self.update_text)
            self.thread.start()

    def stop_video(self):
        if self.thread is not None and self.thread.isRunning():
            self.thread.stop()
            self.thread = None

    def update_image(self, image):
        self.video_frame.setPixmap(QPixmap.fromImage(image))

    def update_text(self, text):
        self.text_area.setText(text)

    def save_text(self):
        response = QMessageBox.question(self, '保存', '确认保存识别的文本吗？', QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        if response == QMessageBox.Yes:
            with open("recognized_text.txt", "w") as file:
                file.write(self.text_area.toPlainText())
            QMessageBox.information(self, '保存成功', '文本已成功保存。')

if __name__ == '__main__':
    app = QApplication(sys.argv)
    mainWin = MainWindow()
    mainWin.show()
    sys.exit(app.exec_())